"""
    Actor and Critic artificial neural networks (ANNs) for the DDPG algorithm.
"""
import tensorflow as tf

from tensorflow.keras.models import load_model, model_from_json
from tensorflow.keras.layers import Input, Dense, Flatten, LeakyReLU
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.regularizers import L2
from tensorflow.keras.models import Model as tfModel

regularizer = L2(2.5e-8)
initializer = RandomNormal(stddev=0.1)

def organizor(num_states, lr: float = 1e-3, feature_num=128) -> tf.keras.Model:
    """
    Actor ANN for the DDPG algorithm, using batch normalization layers.

    Args:
        num_states (int): Number of states.
        num_actions (int): Number of actions.
        lr (float, optional): Learning rate for the optimizer.
            Defaults to 1e-3.

    Returns:
        tf.keras.Model: Actor ANN.
    """

    input_layer = Input(
        shape=num_states
        )
    
    layer1a = Dense(
        feature_num,
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        )(input_layer)

    layer1b_ = Dense(
        feature_num * 4,
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="relu"
        )(input_layer)

    layer1b = Dense(
        feature_num,
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(layer1b_)

    layer1 = layer1a * layer1b

    layer2 = Dense(
        feature_num, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="relu"
        )(layer1)

    output_layer = Dense(
        1, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        activation="sigmoid",
        )(layer2)
    
    model = tf.keras.Model(input_layer, output_layer)
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=lr))

    return model

def actor(num_states: int, num_actions: int, lr: float = 1e-3,
          feature_num=128) -> tf.keras.Model:
    """
    Actor ANN for the DDPG algorithm, using batch normalization layers.

    Args:
        num_states (int): Number of states.
        num_actions (int): Number of actions.
        lr (float, optional): Learning rate for the optimizer.
            Defaults to 1e-3.

    Returns:
        tf.keras.Model: Actor ANN.
    """

    input_layer = Input(
        shape=num_states
        )

    layer1a = Dense(
        feature_num * max(1, round(16/num_states[1])),
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        )(input_layer)

    layer1b_ = Dense(
        feature_num * max(1, round(64/num_states[1])),
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="relu"
        )(input_layer)

    layer1b = Dense(
        feature_num * max(1, round(16/num_states[1])),
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(layer1b_)

    layer1 = tf.reduce_mean(layer1a * layer1b, axis=-2)
    
    layer2a = Dense(
        feature_num, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        )(layer1)

    layer2b_ = Dense(
        feature_num * 4, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(layer1)

    layer2b = Dense(
        feature_num, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(layer2b_)

    layer2 = layer2a * layer2b

    output_layer_ = Dense(
        num_actions * 4, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        activation="relu",
        )(layer2)

    output_layer = Dense(
        num_actions, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        activation="tanh",
        )(output_layer_)

    model = tf.keras.Model(input_layer, output_layer)
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=lr))

    return model


def critic(num_states: int, num_actions: int, lr: float = 1e-3,
           feature_num=128) -> tf.keras.Model:
    """
    Critic ANN for the DDPG algorithm, using batch normalization layers.

    Args:
        num_states (int): Number of states.
        num_actions (int): Number of actions.
        lr (float, optional): Learning rate for the optimizer.
            Defaults to 1e-3.

    Returns:
        tf.keras.Model: Critic ANN.
    """

    # State as input
    state_input = Input(
        shape=num_states
        )

    state_layer1a = Dense(
        feature_num * max(1, round(16/num_states[1])),
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        )(state_input)

    state_layer1b_ = Dense(
        feature_num * max(1, round(64/num_states[1])),
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="relu"
        )(state_input)

    state_layer1b = Dense(
        feature_num * max(1, round(16/num_states[1])),
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(state_layer1b_)

    state_layer1 = tf.reduce_mean(state_layer1a * state_layer1b, axis=-2)
    
    state_layer2a = Dense(
        feature_num, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        )(state_layer1)

    state_layer2b_ = Dense(
        feature_num * 4, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(state_layer1)

    state_layer2b = Dense(
        feature_num, 
        use_bias=True,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        bias_initializer=initializer,
        bias_regularizer=regularizer,
        activation="sigmoid"
        )(state_layer2b_)

    state_layer2 = state_layer2a * state_layer2b

    state_output_layer_ = Dense(
        num_actions * 4, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        activation="relu",
        )(state_layer2)

    state_act = Dense(
        num_actions, 
        use_bias=False,
        kernel_initializer=initializer,
        kernel_regularizer=regularizer,
        activation="tanh",
        )(state_output_layer_)

    # Action as input
    action_input = tf.keras.layers.Input(shape=(num_actions))
    action_dense = tf.keras.layers.Dense(num_actions*4,
                                         activation="relu",
                                         use_bias=True,
                                         kernel_initializer=initializer,
                                         kernel_regularizer=regularizer,
                                         bias_initializer=initializer,
                                         bias_regularizer=regularizer,
                                         )(action_input)
    action_act = tf.keras.layers.Dense(num_actions,
                                       activation="tanh",
                                       use_bias=True,
                                       kernel_initializer=initializer,
                                       kernel_regularizer=regularizer,
                                       bias_initializer=initializer,
                                       bias_regularizer=regularizer,
                                       )(action_dense)

    concat = tf.keras.layers.Concatenate()([state_act, action_act])

    dense1 = tf.keras.layers.Dense(feature_num*4,
                                   activation="relu",
                                   use_bias=True,
                                   kernel_initializer=initializer,
                                   kernel_regularizer=regularizer,
                                   bias_initializer=initializer,
                                   bias_regularizer=regularizer,
                                   )(concat)
    dense2 = tf.keras.layers.Dense(feature_num,
                                   activation="relu",
                                   use_bias=True,
                                   kernel_initializer=initializer,
                                   kernel_regularizer=regularizer,
                                   bias_initializer=initializer,
                                   bias_regularizer=regularizer,
                                   )(dense1)
    last_init = tf.random_uniform_initializer(minval=-0.003, maxval=0.003)
    outputs = tf.keras.layers.Dense(1,
                                    use_bias=True,
                                    activation="tanh",
                                    kernel_initializer=initializer,
                                    kernel_regularizer=regularizer,
                                    bias_initializer=initializer,
                                    bias_regularizer=regularizer,
                                    )(dense2)

    model = tf.keras.Model([state_input, action_input], outputs)
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=lr))

    return model
